from openai import OpenAI  # 确保安装 openai>=1.0
from typing import List, Dict

# 初始化客户端
client = OpenAI(
    base_url="http://127.0.0.1:8102/v1",  # 替换为实际地址
    api_key="EMPTY"
)

# 对话历史存储
chat_history: List[Dict] = []

while True:
    # 用户输入
    user_input = input("\n用户: ")
    if user_input.lower() in ["exit", "quit"]:
        break

    # 添加用户消息到历史
    chat_history.append({"role": "user", "content": user_input})

    # 调用模型 (使用chat.completions接口)
    response = client.chat.completions.create(
        model="/data/model/models/ZhipuAI/chatglm3-6b",  # 需与vLLM启动时的模型名一致
        messages=chat_history,  # 传入完整历史
        max_tokens=500,
        temperature=0.7,
        stream=False
    )

    # 获取模型回复
    ai_response = response.choices[0].message.content
    print(f"\nAI: {ai_response}")

    # 添加AI回复到历史
    chat_history.append({"role": "assistant", "content": ai_response})

    # 控制历史长度（防止内存溢出）
    if len(chat_history) > 10:  # 保留最近5轮对话（每轮2条消息）
        chat_history = chat_history[-10:]